import java.util.*;

/**
* @description: 863. 二叉树中所有距离为 K 的结点
* @author hewei
* @date 2022/9/9 15:20
* @version 1.0
*/

public class DistanceK {

    HashMap<TreeNode, TreeNode> map = new HashMap<>();
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        process(root);
        Deque<TreeNode> queue = new LinkedList<>();
        HashSet<TreeNode> set = new HashSet<>();
        set.add(target);
        queue.add(target);
        List<Integer> ans = new ArrayList<>();
        if (k == 0) {
            ans.add(target.val);
            return ans;
        }
        while (k-- > 0) {
            int size = queue.size();
            while (size-- > 0) {
                TreeNode poll = queue.poll();
                if (poll.left != null && !set.contains(poll.left)) {
                    queue.add(poll.left);
                    set.add(poll.left);
                    if (k == 0) ans.add(poll.left.val);
                }
                if (poll.right != null && !set.contains(poll.right)) {
                    queue.add(poll.right);
                    set.add(poll.right);
                    if (k == 0) ans.add(poll.right.val);
                }
                if (map.containsKey(poll) && !set.contains(map.get(poll))) {
                    queue.add(map.get(poll));
                    set.add(map.get(poll));
                    if (k == 0) ans.add(map.get(poll).val);
                }
            }
        }
        return ans;
    }

    public void process(TreeNode node) {
        if (node == null || (node.left == null && node.right == null)) return;

        if (node.left != null) {
            map.put(node.left, node);
            process(node.left);
        }
        if (node.right != null) {
            map.put(node.right, node);
            process(node.right);
        }
    }
}
